package cas.ihep.spark.loop;

import cas.ihep.spark.io.Packet;
import io.netty.bootstrap.Bootstrap;
import io.netty.buffer.Unpooled;
import io.netty.channel.*;
import io.netty.channel.nio.NioEventLoopGroup;
import io.netty.channel.socket.DatagramPacket;
import io.netty.channel.socket.nio.NioDatagramChannel;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.net.SocketAddress;
import java.nio.ByteBuffer;
import java.util.concurrent.ConcurrentHashMap;

public class NettyUDPClient extends IterationNetClient {
    private Bootstrap b;
    private SocketAddress localAddr;
    NettyUDPClient(SocketAddress addr){
        b = new Bootstrap();
        localAddr=addr;
        EventLoopGroup group = new NioEventLoopGroup();
        b.group(group)
                .channel(NioDatagramChannel.class)
                .option(ChannelOption.SO_BROADCAST, true);
    }

    @Override
    public void waitForPacket(OnPacketRead onRead, int timeOut) {
        try{
            Handler h=new Handler(onRead);
            ChannelFuture chn=b.handler(h).bind(localAddr);
            chn.sync().channel().closeFuture().await();chn.channel();
        } catch (InterruptedException e) {
            e.printStackTrace();
        }
    }

    private static class Handler extends SimpleChannelInboundHandler<DatagramPacket> {
        OnPacketRead listener;
        ConcurrentHashMap<Integer,Packet> objCache;
        Handler(OnPacketRead onRead){
            listener=onRead;
            objCache=new ConcurrentHashMap<>();
        }
        @Override
        protected void channelRead0(ChannelHandlerContext ctx, DatagramPacket dp) throws Exception {
            Packet packet=Packet.wrap(dp.content().nioBuffer());
            int tp=packet.type();
            boolean isRunning;
            InetSocketAddress from=dp.sender();
            VirtualNetContext vtx=new VirtualNetContext() {
                @Override
                public void write(Packet packet, SocketAddress to) throws IOException {
                    if(packet.type()==Packet.DataReply) {
                        objCache.put(packet.numOfIterations(), packet);
                    }
                    ctx.writeAndFlush(new DatagramPacket(Unpooled.wrappedBuffer(packet.getBuffer()),(InetSocketAddress)to));
                }
            };
            if(tp==Packet.DataRequest){
                int iternum=packet.numOfIterations();
                Packet result=objCache.get(iternum);
                if(result==null) {
                    isRunning = listener.onRead(vtx, packet, from);
                }else{
                    ByteBuffer ibuffer=result.getBuffer();
                    ctx.writeAndFlush(new DatagramPacket(Unpooled.wrappedBuffer(ibuffer),from));
                    //ibuffer.rewind();
                    isRunning=true;
                }
            }else{
                isRunning=listener.onRead(vtx,packet,from);
            }
            if(!isRunning){
                ctx.close();
            }
        }
    }
}
